"""
工具函数 - ResNeXt 图像分类训练监控平台
"""
import os
import json
import torch
import torch.nn as nn
from torchvision import models, transforms
from datetime import datetime
import glob

def get_device():
    """获取最佳的计算设备"""
    if torch.cuda.is_available():
        return torch.device("cuda")
    elif torch.backends.mps.is_available():   
        return torch.device("mps")
    else:
        return torch.device("cpu")

def get_model_transform(is_training=True):
    """获取标准化的图像预处理transforms"""
    if is_training:
        return transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])
    else:
        return transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])

def create_resnext_model(num_classes, device):
    """创建并初始化ResNeXt模型"""
    try:
        # 使用新的权重参数名称
        model = models.resnext50_32x4d(weights=models.ResNeXt50_32X4D_Weights.DEFAULT)
        num_ftrs = model.fc.in_features
        model.fc = nn.Linear(num_ftrs, num_classes)
        model = model.to(device)
        return model
    except Exception as e:
        # 兼容旧版本PyTorch
        print(f"Warning: {e}. Falling back to legacy weight loading.")
        model = models.resnext50_32x4d(pretrained=True)
        num_ftrs = model.fc.in_features
        model.fc = nn.Linear(num_ftrs, num_classes)
        model = model.to(device)
        return model

def cleanup_gpu_memory(device):
    """清理GPU内存"""
    if device.type == 'cuda':
        torch.cuda.empty_cache()

def get_latest_model_file(model_folder):
    """查找最新的模型文件"""
    try:
        model_files = glob.glob(os.path.join(model_folder, 'trained_model_*.pth'))
        if not model_files:
            return None
        latest_file = max(model_files, key=os.path.getctime)
        return os.path.basename(latest_file)
    except Exception as e:
        print(f"Error getting latest model: {e}")
        return None

def validate_json_structure(json_data):
    """验证JSON数据结构"""
    if not isinstance(json_data, list):
        return False, "JSON数据应该是一个列表"
    
    for i, item in enumerate(json_data):
        if 'image' not in item:
            return False, f"第{i+1}个项目缺少'image'字段"
        
        image_data = item['image']
        if 'url' not in image_data or 'point_code' not in image_data:
            return False, f"第{i+1}个项目的image字段缺少'url'或'point_code'"
    
    return True, "JSON结构验证通过"

def safe_filename(filename):
    """生成安全的文件名"""
    import re
    # 移除或替换不安全的字符
    filename = re.sub(r'[<>:"/\\|?*]', '', filename)
    return filename

def format_file_size(size_bytes):
    """格式化文件大小显示"""
    if size_bytes < 1024:
        return f"{size_bytes} B"
    elif size_bytes < 1024 * 1024:
        return f"{size_bytes / 1024:.1f} KB"
    elif size_bytes < 1024 * 1024 * 1024:
        return f"{size_bytes / (1024 * 1024):.1f} MB"
    else:
        return f"{size_bytes / (1024 * 1024 * 1024):.1f} GB"
